// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/benchmark.h>
#include <melon/init/init.h>
#include <turbo/flags/flag.h>

#include <pollux/common/base/compare_flags.h>
#include <pollux/functions/lib/benchmarks/FunctionBenchmarkBase.h>
#include <pollux/vector/vector_fuzzer.h>

TURBO_FLAG(int64_t,fuzzer_seed, 99887766, "Seed for random input dataset generator");

using namespace kumo::pollux;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::test;

namespace {

class VectorCompareBenchmark : public functions::test::FunctionBenchmarkBase {
 public:
  explicit VectorCompareBenchmark(size_t vectorSize)
      : FunctionBenchmarkBase(), vectorSize_(vectorSize), rows_(vectorSize) {
    VectorFuzzer::Options opts;
    opts.vectorSize = vectorSize_;
    opts.nullRatio = 0;
    opts.containerVariableLength = 1000;
    VectorFuzzer fuzzer(opts, pool(), FLAGS_fuzzer_seed);

    flatVector_ = fuzzer.fuzzFlat(BIGINT());

    arrayVector_ = fuzzer.fuzzFlat(ARRAY(BIGINT()));

    mapVector_ = fuzzer.fuzzFlat(MAP(BIGINT(), BIGINT()));

    rowVector_ = fuzzer.fuzzFlat(ROW({BIGINT(), BIGINT(), BIGINT()}));
  }

  size_t run(const VectorPtr& vector) {
    size_t sum = 0;

    for (auto i = 0; i < vectorSize_; i++) {
      sum += *vector->compare(vector.get(), i, i, kFlags);
    }
    melon::doNotOptimizeAway(sum);
    return vectorSize_;
  }

  // Avoid dynamic dispatch by casting the vector before calling compare to its
  // derived that have final compare function.
  size_t runFastFlat() {
    size_t sum = 0;
    auto flat_vector = flatVector_->as<FlatVector<int64_t>>();
    for (auto i = 0; i < vectorSize_; i++) {
      sum += *flat_vector->compare(flat_vector, i, vectorSize_ - i - 1, kFlags);
    }
    melon::doNotOptimizeAway(sum);
    return vectorSize_;
  }

  VectorPtr flatVector_;
  VectorPtr arrayVector_;
  VectorPtr mapVector_;
  VectorPtr rowVector_;

 private:
  static constexpr CompareFlags kFlags{
      true,
      true,
      false,
      CompareFlags::NullHandlingMode::kNullAsValue};

  const size_t vectorSize_;
  SelectivityVector rows_;
};

std::unique_ptr<VectorCompareBenchmark> benchmark;

BENCHMARK(compareSimilarSimpleFlat) {
  benchmark->run(benchmark->flatVector_);
}

BENCHMARK(compareSimilarSimpleFlatNoDispatch) {
  benchmark->runFastFlat();
}

BENCHMARK(compareSimilarArray) {
  benchmark->run(benchmark->arrayVector_);
}

BENCHMARK(compareSimilarMap) {
  benchmark->run(benchmark->mapVector_);
}

BENCHMARK(compareSimilarRow) {
  benchmark->run(benchmark->rowVector_);
}

BENCHMARK_DRAW_LINE();

} // namespace

int main(int argc, char* argv[]) {
  melon::Init init{&argc, &argv};
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  memory::MemoryManager::initialize({});
  benchmark = std::make_unique<VectorCompareBenchmark>(1000);
  melon::runBenchmarks();
  benchmark.reset();
  return 0;
}
